%-----------------> REWRITE METHOD BODIES
change_methodbodycount(BC, NewBC) :-
	NewBC is BC.

change_methodbodyinfo(MBs, MarkedList, CP, NewMBs, NewMarkedList) :-

	empty_assoc(OTN),
	empty_assoc(NTO),
	b_setval(oldtonew,OTN),
	b_setval(newtoold,NTO),

	empty_assoc(MKD),
	b_setval(markedinstr,MKD),

	b_setval(jumptable,[]),

	%spy(modify_method_body),

	add_sec_var_instructions(MBs,CP,MarkedList,TempMBs,TempMarkedList),

	fix_them_branches(TempMBs, NewMBs),

	fix_them_marks(NewMBs, TempMarkedList, NewMarkedList).

fix_them_marks([], [], []).
fix_them_marks([MB|MBs], [(Method-Marks)|Rest], [(Method-NewMarks)|NewRest]) :-
	MB =.. [method_body, Method, _MS, _LC, _ISD, _MSD, _CL, C, _EC, _E, _TC, _T],
	fix_that_there_marking(C, Marks, NewMarks),
	fix_them_marks(MBs, Rest, NewRest).

fix_that_there_marking([], [], []).
fix_that_there_marking([I|Is], [Mark|Marks], [(I,Mark)|NewMarks]) :-
	fix_that_there_marking(Is, Marks, NewMarks).

fix_them_branches([],[]).
fix_them_branches([TMB|TM], [NMB|NM]) :-

	TMB =.. [method_body, MethodNum, MS, LC, ISD, MSD, CL, C, EC, E, TC, T],

	fix_that_there_method(MethodNum, C, NewC),

	NMB =.. [method_body, MethodNum, MS, LC, ISD, MSD, CL, NewC, EC, E, TC, T],

	fix_them_branches(TM, NM).

% Base Case
fix_that_there_method(_,[],[]).

% Not a Branch Case
fix_that_there_method( MN, [(Offset-H)|TRest], [(Offset-H)|NRest]) :-
	H =.. [Head | _Args],
	\+isABranch(Head),
	fix_that_there_method(MN, TRest, NRest).

% LookUpSwitch Case
fix_that_there_method(MN, [(Offset-H)|TRest], [(Offset-NewH)|NRest]) :-
	H =.. [lookupswitch, DefaultOffset, CaseCount, CaseOffsets],

	b_getval(markedinstr, MKD),

	((get_assoc((MN-Offset), MKD, _VAL)) ->
		NewH =.. [lookupswitch, DefaultOffset, CaseCount, CaseOffsets],
		fix_that_there_method(MN, TRest, NRest)
	;
		b_getval(oldtonew, OTN),
		b_getval(newtoold, NTO),

		% Get the Original Offset
		get_assoc((MN-Offset), NTO, (MN-OldOffset)),
		OldDefaultTarget is (OldOffset+DefaultOffset),
		get_assoc((MN-OldDefaultTarget), OTN, (MN-NewDefaultTarget)),
		NewDefaultOffset is (NewDefaultTarget-Offset),

		fix_case_offsets(MN, Offset, OldOffset, CaseOffsets, NewCaseOffsets),

		%	writeln(CaseOffsets),
		%	writeln(NewCaseOffsets),

		NewH =.. [lookupswitch, NewDefaultOffset, CaseCount, NewCaseOffsets],
		fix_that_there_method(MN, TRest, NRest)
	).

% Branch or Jump Case
fix_that_there_method(MN, [(Offset-H)|TRest], [(Offset-NewH)|NRest]) :-
	H =.. [Head | Tail], 
	Tail = [OldJumpOffset],

	

	b_getval(markedinstr, MKD),
	((get_assoc((MN-Offset), MKD, _VAL)) ->
		b_getval(oldtonew, OTN),
		b_getval(newtoold, NTO),

		get_assoc((MN-OldJumpOffset), OTN, (MN-NewTarget)),
		NewJumpOffset is (NewTarget-(Offset+4)),

		NewTail = [NewJumpOffset],
		NewH =.. [Head | NewTail],
        
		fix_that_there_method(MN, TRest, NRest)
	;
		b_getval(oldtonew, OTN),
		b_getval(newtoold, NTO),

		get_assoc((MN-Offset), NTO, (MN-OldOffset)),
		OldTarget is (OldOffset+4+OldJumpOffset),
		get_assoc((MN-OldTarget), OTN, (MN-NewTarget)),
		NewJumpOffset is (NewTarget-(Offset+4)),

		NewTail = [NewJumpOffset],
		NewH =.. [Head | NewTail],
        
		fix_that_there_method(MN, TRest, NRest)
	).


% Helper function, determines if the instruction is a branching jump instruction
isABranch(ifeq).
isABranch(iffalse).
isABranch(ifge).
isABranch(ifgt).
isABranch(ifle).
isABranch(iflt).
isABranch(ifnge).
isABranch(ifngt).
isABranch(ifnle).
isABranch(ifnlt).
isABranch(ifne).
isABranch(ifstricteq).
isABranch(ifstrictne).
isABranch(iftrue).
isABranch(jump).
isABranch(lookupswitch).
	
%=====================================================================================
add_sec_var_instructions(MBs, CP, MarkedList, NewMBs, NewMarkedList) :- 

	find_index_multinameCP('navigateToURL', CP, InstrIndex),

	b_setval(instrindices,[InstrIndex]),
	
	find_index_multinameCP('nativeApplication', CP, NativeAppIndex),
	b_setval(nativeappindex,NativeAppIndex),
	find_index_multinameCP('exit', CP, HaltIndex),
	b_setval(haltindex,HaltIndex),
	find_index_multinameCP('isOkay',CP,SecVarIndex),

	insert_guards_overall(MBs, InstrIndex, SecVarIndex, HaltIndex, NativeAppIndex, MarkedList, NewMBs, NewMarkedList).
	

%----------------------------------------------------------

%Search through each method body to see if it has the security 
%relevant instruction given by the method name residing at 
%MultinameIndex. Insert guards around that instruction.
%For right now, assuming that the sec rel event is going to be a 
%method call, of the form callpropvoid(index, arg_count). 
%TODO: change insert_instr_method_body to handle other method
%calls such as callproplex(index, arg_count), or callstatic(i, a).
%
%
%TODO: Not sure if setlocal_1 and getlocal_1 are okay to use 
%here. Could there be other values stored in those registers?
%How do we know whether they're free to use?
%
%InstrINdex is the multiname CP index of the security relevant instruction.
%SecVarIndex is the multiname CP index of the security variable.
%HaltIndex is the multiname CP index of the halt method.

insert_guards_overall([], _InstrIndex, _SecVarIndex, _HaltIndex, _NativeAppIndex, MarkedList, [], MarkedList).

insert_guards_overall([MB|MBs], InstrIndex, SecVarIndex, HaltIndex, NativeAppIndex, [(Method-MarkedC)|MarkedCs], 
					[NewMB|NewMBs], [(Method-NewMarkedC)|NewMarkedCs]) :- 

	MB =..  [method_body, Method, MS, LC, ISD, MSD, CL, C, EC, E, TC, T],
	
	InsertInstructions = [],
	NewLC is LC+1,

	modify_method_body(C, NewC, MarkedC, NewMarkedC, 0, NewCL, NewLC, Method, 0, [], InsertInstructions),

	NewLC is LC + 1,
	NewMS is MS + 2,

	NewMB =.. [method_body, Method, NewMS, NewLC, ISD, MSD, NewCL, NewC, EC, E, TC, T],

	insert_guards_overall(MBs, InstrIndex, SecVarIndex, HaltIndex, NativeAppIndex, MarkedCs, NewMBs, NewMarkedCs).

print_my_code([]).
print_my_code([C|Cs]) :-
	writeln(C),
	print_my_code(Cs).

print_my_methods([]).
print_my_methods([C|Cs]) :-
	C =.. [method_body, Method, NewMS, NewLC, ISD, MSD, NewCL, NewC, EC, E, TC, T],
	write(Method), writeln(': '),
	print_my_method(NewC),
	writeln(' '),
	print_my_methods(Cs).

print_my_method([]).
print_my_method([C|Cs]) :-
	writeln(C),
	print_my_method(Cs).


create_insert_list(_, _, _, CaseCount, CaseCount, [], _, _, FIL, FIL, FML, FML).
create_insert_list(MethodNum, Offset, BaseAdd, CaseCount, FullCaseCount, [CO | CORest], CasesOffset, LC, IIL, FIL, IML, FML) :-
	NextCaseCount is CaseCount+1,

	instruction_size(getlocal(0), GLLength),
	instruction_size(pushshort(0), PushLength),
	instruction_size(iflt(0), BranchLength),

	GLOffset is CasesOffset,
	PushOffset is GLOffset+GLLength,
	BranchOffset is PushOffset+PushLength,
	NewCasesOffset is BranchOffset+BranchLength,

	SumAdd is BaseAdd + (GLLength+PushLength+BranchLength),
	COOffset is Offset + CO,

	b_getval(markedinstr,MKD1),

	put_assoc((MethodNum-BranchOffset), MKD1, 1, MKD2),

	b_setval(markedinstr,MKD2),

	append(IIL, [(GLOffset-getlocal(LC)), (PushOffset-pushshort(CaseCount)), (BranchOffset-ifeq(COOffset))], IIL1),
	append(IML, ['Marked', 'Marked', 'Marked'], IML1),

	create_insert_list(MethodNum, Offset, SumAdd, NextCaseCount, FullCaseCount, CORest, NewCasesOffset, LC, IIL1, FIL, IML1, FML).

%-----------------------------------------------

% modify_method_body(InsList, NewInsList, MarkedInsList, NewMarkedInsLIst, 

modify_method_body([], [], [], [], CL, CL, _, _MethodNum, _CurrentOffset, _MultinameIndexes, _InsertList).

modify_method_body([Offset - I | Rest], NewInsList, [MarkIns | MarkedInsList], NewMarkedList, 
		   CL, NewCL, LC, MethodNum, CurrentOffset, MultinameIndexes, InsertList) :-

			   %((MethodNum = 18) -> spy(create_insert_list) ; true),

	I =.. [lookupswitch|Tail], !,
	Tail = [DefaultOffset, CaseCount, CaseOffsets],

	instruction_size(setlocal(0), SLLength),
	instruction_size(getlocal(0), GLLength),
	instruction_size(pushshort(0), PushLength),
	instruction_size(iflt(0), BranchLength),

	instruction_size(I, CLInstr),

	SL1 is CurrentOffset,
	GL1 is SL1+SLLength,
	PS1 is GL1+GLLength,
	IF1 is PS1+PushLength,
	GL2 is IF1+BranchLength,
	PS2 is GL2+GLLength,
	IF2 is PS2+PushLength,
	CasesOffsets is IF2+BranchLength,

	BaseAdd is (SLLength + GLLength + PushLength + BranchLength + GLLength + PushLength + BranchLength),

	DO1 is (Offset + DefaultOffset),

	b_getval(markedinstr,MKD1),

	put_assoc((MethodNum-IF1), MKD1, 1, MKD2),
	put_assoc((MethodNum-IF2), MKD2, 1, MKD3),

	b_setval(markedinstr,MKD3),

	InitialInsertList = [(SL1-setlocal(LC)),
			     (GL1-getlocal(LC)),
			     (PS1-pushshort(0)),
			     (IF1-iflt(DO1)),
			     (GL2-getlocal(LC)),
			     (PS2-pushshort(CaseCount)),
			     (IF2-ifgt(DO1))],

	InitialMarkList = ['Marked', 'Marked', 'Marked', 'Marked', 'Marked', 'Marked', 'Marked'],

	NextCaseCount is CaseCount+1,
	create_insert_list(MethodNum, Offset, BaseAdd, 0, NextCaseCount, CaseOffsets, CasesOffsets, LC, 
			   InitialInsertList, FullInsertList, InitialMarkList, FullMarkList),
	
	% Add the offsets to the new assoc list
	b_getval(oldtonew,OTN),
	b_getval(newtoold,NTO),

	put_assoc((MethodNum-Offset), OTN, (MethodNum-CurrentOffset), OTN2),
	put_assoc((MethodNum-CurrentOffset), NTO, (MethodNum-Offset), NTO2),

	b_setval(oldtonew,OTN2),
	b_setval(newtoold,NTO2),

	CaseMultiplier is (GLLength+PushLength+BranchLength),
	FullCaseCount is CaseCount+1,
	SumSize is BaseAdd + (CaseMultiplier*FullCaseCount),

	% Increment the instruction count
	TempCL is (CL+SumSize),

	% Find out the length of the old Instr	
	NewCurOffset is (CurrentOffset + SumSize),

	% Recursively call
	modify_method_body(Rest, NewRest, MarkedInsList, NewMarkedInsList, TempCL, NewCL, LC, 
			   MethodNum, NewCurOffset, MultinameIndexes, InsertList),

	append(FullMarkList, NewMarkedInsList, NewMarkedList),
	append(FullInsertList, NewRest, NewInsList).

modify_method_body([Offset - I | Rest], [CurrentOffset - I | NewRest], 
		   [Marking | MarkList], [NewMarking | NewMarkList], 
		   CL, NewCL, LC, MethodNum, CurrentOffset, MultinameIndexes, InsertList) :-

	% Add the offsets to the new assoc list
	b_getval(oldtonew,OTN),
	b_getval(newtoold,NTO),

	put_assoc((MethodNum-Offset), OTN, (MethodNum-CurrentOffset), OTN2),
	put_assoc((MethodNum-CurrentOffset), NTO, (MethodNum-Offset), NTO2),

	b_setval(oldtonew,OTN2),
	b_setval(newtoold,NTO2),

	% For this policy, all original instructions are left unmarked
	NewMarking = Marking, 

	% Find out the length of the old Instr	
	instruction_size(I, CLInstr),
	NextCurOffset is (CurrentOffset + CLInstr),

	% Add 1 to the instruction count
	TempCL is (CL+CLInstr),

	% Recursively call
	modify_method_body(Rest, NewRest, MarkList, NewMarkList, TempCL, NewCL, LC, 
			   MethodNum, NextCurOffset, MultinameIndexes, InsertList). 

is_jump(ifeq).
is_jump(iffalse).
is_jump(ifge).
is_jump(ifgt).
is_jump(ifle).
is_jump(iflt).
is_jump(ifnge).
is_jump(ifngt).
is_jump(ifnle).
is_jump(ifnlt).
is_jump(ifne).
is_jump(ifnge).
is_jump(ifstricteq).
is_jump(ifstrictne).
is_jump(iftrue).
is_jump(jump).


